import gym
from gym import spaces
import pyautogui
import cv2
import numpy as np
import time
import os
from PIL import Image
import torch
from transformers import AutoImageProcessor, AutoModel

class NotepadEnv(gym.Env):
    def __init__(self, render_mode=None):
        super(NotepadEnv, self).__init__()
        
        self.action_space = spaces.Discrete(4)
        
        # 加载DINO模型
        pretrained_model_name = os.path.abspath("E:\\code\\ppo_text\\model\\dinov3-convnext-tiny-pretrain-lvd1689m")
        self.processor = AutoImageProcessor.from_pretrained(pretrained_model_name)
        self.dino_model = AutoModel.from_pretrained(pretrained_model_name, device_map="auto")
        self.dino_model.cuda()
        
        # 获取DINO特征维度 (根据模型输出确定)
        dummy_image = Image.new('RGB', (224, 224))
        dummy_inputs = self.processor(images=dummy_image, return_tensors="pt")
        dummy_inputs['pixel_values'] = dummy_inputs['pixel_values'].to('cuda')
        with torch.no_grad():
            dummy_outputs = self.dino_model(dummy_inputs['pixel_values'])
            feature_dim = dummy_outputs.last_hidden_state.shape[-1]  # 通常是768或1024
        
        # 更新观察空间为DINO特征维度
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(feature_dim,), dtype=np.float32
        )
        
        # 使用相对目录路径
        self.text_dir = os.path.join(os.getcwd(), "text")
        os.makedirs(self.text_dir, exist_ok=True)  # 确保目录存在
        
        self.step_count = 0
        self.max_steps = 20
        self.episode_count = 0  # 记录episode数量
        
        # 用于累积奖励的变量
        self.text_entered = False
        self.file_saved = False
        self.notepad_closed = False
        
        # 存储每轮的奖励信息
        self.episode_rewards = []
        self.current_episode_reward = 0
        
        # 动作名称映射
        self.action_names = {
            0: "打开记事本",
            1: "输入文本",
            2: "保存文件",
            3: "关闭文件"
        }

        # 添加状态跟踪变量
        self.notepad_opened = False
        self.notepad_open_count = 0  # 记录打开记事本的次数
        self.text_enter_count = 0     # 记录输入文本的次数
        self.last_action = None
        self.repeated_action_count = 0  # 用于追踪重复动作次数
        
        # 添加中间步骤奖励标记
        self._text_enter_rewarded = False
        self._file_save_rewarded = False
        
        # 奖励归一化相关变量
        self.reward_history = []  # 存储历史奖励用于计算均值和标准差
        self.reward_mean = 0.0
        self.reward_std = 1.0
        self.reward_history_size = 1000  # 保留最近1000个奖励值
        
        # 奖励缩放因子
        self.reward_scale = 0.5  # 降低奖励缩放因子

    def get_obs(self):
        """获取当前屏幕截图并通过DINO提取特征"""
        # 获取屏幕截图
        screenshot = pyautogui.screenshot()
        frame = np.array(screenshot)
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        
        # 调整图像大小以适应DINO输入要求
        frame = cv2.resize(frame, (224, 224))  # DINO通常使用224x224
        
        # 转换为PIL图像
        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        
        # 使用DINO processor处理图像
        inputs = self.processor(images=pil_image, return_tensors="pt")
        inputs['pixel_values'] = inputs['pixel_values'].to('cuda')
        
        # 提取特征
        with torch.no_grad():
            outputs = self.dino_model(inputs['pixel_values'])
            # 使用[CLS] token作为特征表示
            features = outputs.last_hidden_state[:, 0, :].cpu().numpy()
            
        return features.flatten().astype(np.float32)

    def reset(self):
        # 如果不是第一次调用reset，说明上一轮已经结束
        if self.step_count > 0:
            self.episode_count += 1
            self.episode_rewards.append(self.current_episode_reward)
            
            # 输出本轮奖励信息
            print(f"Episode {self.episode_count} 结束，总奖励: {self.current_episode_reward}, 步数: {self.step_count}")
            
            # 重置当前episode奖励
            self.current_episode_reward = 0
        
        self.step_count = 0
        self.text_entered = False
        self.file_saved = False
        self.notepad_closed = False
        self.notepad_opened = False
        self.notepad_open_count = 0  # 重置打开记事本次数
        self.text_enter_count = 0    # 重置输入文本次数
        self.task_reward_given = False  # 标记是否已给予任务完成奖励
        self.last_action = None
        self.repeated_action_count = 0  # 重置重复动作计数
        
        # 重置中间步骤奖励标记
        self._text_enter_rewarded = False
        self._file_save_rewarded = False
        
        # 删除text目录下所有txt文件
        try:
            txt_files = [f for f in os.listdir(self.text_dir) if f.endswith('.txt')]
            for txt_file in txt_files:
                file_path = os.path.join(self.text_dir, txt_file)
                try:
                    os.remove(file_path)
                except:
                    pass
        except:
            pass
        
        # 关闭记事本窗口
        self.close_notepad_windows()
            
        print(f"Episode {self.episode_count + 1} 开始")
        return self.get_obs()

    def normalize_reward(self, reward):
        """奖励归一化函数"""
        # 将当前奖励添加到历史记录中
        self.reward_history.append(reward)
        
        # 保持历史记录大小
        if len(self.reward_history) > self.reward_history_size:
            self.reward_history.pop(0)
        
        # 计算均值和标准差
        if len(self.reward_history) > 1:
            self.reward_mean = np.mean(self.reward_history)
            self.reward_std = np.std(self.reward_history)
            
            # 避免除零错误
            if self.reward_std > 1e-8:
                # 归一化奖励: (reward - mean) / (std + epsilon)
                normalized_reward = (reward - self.reward_mean) / (self.reward_std + 1e-8)
                # 应用缩放因子
                normalized_reward *= self.reward_scale
                return normalized_reward
        
        # 如果历史记录不足或标准差为0，则返回原始奖励
        return reward * self.reward_scale

    def step(self, action):
        # 将numpy数组转换为Python整数
        if isinstance(action, np.ndarray):
            action = int(action.item())
        
        reward = 0
        done = False
        info = {}
        
        self.step_count += 1
        print(f"步骤 {self.step_count}/{self.max_steps}: 执行动作 - {self.action_names.get(action, '未知动作')} (action={action})")

        # 惩罚重复动作（减轻惩罚力度）
        if action == self.last_action:
            self.repeated_action_count += 1
            # 减轻惩罚幅度和增长速度
            penalty = min(2.0, 0.2 * self.repeated_action_count)
            reward -= penalty
            print(f"  -> 轻微惩罚: 重复动作 #{self.repeated_action_count} (-{penalty:.2f})")
        else:
            # 奖励改变动作的行为
            reward += 0.1
            print(f"  -> 奖励: 尝试新动作 (+0.1)")
            self.repeated_action_count = 0  # 重置重复计数

        self.last_action = action

        # 执行动作前确保没有记事本窗口打开（除了保存和关闭操作）
        if action in [0, 1] and not self.is_notepad_running():
            self.close_notepad_windows()  # 确保干净环境
            time.sleep(0.1)

        # 执行动作
        if action == 0:
            # 打开记事本
            self.open_notepad()
            self.notepad_opened = True
            self.notepad_open_count += 1
            
            # 奖励叠加递减机制，但保持最小奖励为0.1
            open_reward = max(0.5, 2.0 - (self.notepad_open_count - 1) * 0.2)
            reward += open_reward
            print(f"  -> 执行打开记事本操作 (第{self.notepad_open_count}次) (+{open_reward:.2f})")
        elif action == 1:
            if self.is_notepad_running():
                pyautogui.write('Hello from RL!', interval=0.1)
                self.text_entered = True
                self.text_enter_count += 1
                
                # 输入文本奖励递减机制，但保持最小奖励为1.0
                text_reward = max(2.0, 5.0 - (self.text_enter_count - 1) * 0.5)
                reward += text_reward
                print(f"  -> 在记事本中输入文本 (第{self.text_enter_count}次) (+{text_reward:.2f})")
                
                # 首次输入文本给予额外奖励
                if not self._text_enter_rewarded:
                    reward += 3.0
                    self._text_enter_rewarded = True
                    print(f"  -> 首次输入文本奖励 (+3.0)")
            else:
                reward -= 1.0  # 减轻惩罚
                print("  -> 记事本未运行，无法输入文本 (-1.0)")
        elif action == 2:  # 保存文件
            if self.is_notepad_running() and self.text_entered:
                # 使用 Ctrl+S 快捷键保存
                pyautogui.hotkey('ctrl', 's')
                time.sleep(0.1)  # 增加等待时间
                
                # 确保焦点在文件名输入框
                try:
                    import pygetwindow as gw
                    save_dialog = gw.getWindowsWithTitle("另存为")
                    if save_dialog:
                        save_dialog[0].activate()
                        time.sleep(0.1)
                except:
                    pass
                
                # 生成唯一的文件名
                filename = f"rl_text_{int(time.time())}"
                pyautogui.write(filename)
                time.sleep(0.1)
                
                # 确认保存
                pyautogui.press('enter')
                time.sleep(1)
                pyautogui.press('enter')
                pyautogui.press('enter')
                self.file_saved = True
                reward += 5.0  # 保存文件的奖励
                print(f"  -> 保存文件为: {filename}.txt (+5.0)")
                
                # 首次保存文件给予额外奖励
                if not self._file_save_rewarded:
                    reward += 3.0
                    self._file_save_rewarded = True
                    print(f"  -> 首次保存文件奖励 (+3.0)")
            else:
                reward -= 1.0  # 减轻惩罚
                print("  -> 记事本未运行或未输入文本，无法保存文件 (-1.0)")

        elif action == 3:  # 关闭文件
            if self.is_notepad_running():
                self.close_notepad_windows()
                self.notepad_closed = True
                # 如果已经完成任务则给予奖励，否则轻微惩罚
                if self.text_entered and self.file_saved:
                    reward += 5.0
                    print("  -> 成功关闭记事本窗口 (+5.0)")
                else:
                    reward -= 1.0  # 减轻惩罚
                    print("  -> 过早关闭记事本 (-1.0)")
            else:
                # 不再惩罚关闭已经关闭的记事本
                print("  -> 记事本未运行，无需关闭 (0.0)")

        # 每一步都检查任务完成情况
        task_completed = self.check_task_completion()
        
        # 如果任务完成，立即结束episode
        if task_completed:
            done = True
            print(f"  -> 任务完成，episode结束")

        # 检查是否达到最大步数
        elif self.step_count >= self.max_steps:
            done = True
            # 减轻未完成任务的惩罚
            if not (self.text_entered and self.file_saved):
                reward -= 2.0  # 从-5.0降低到-2.0
                print(f"  -> 达到最大步数但未完成任务 (-2.0)")
            print(f"  -> 达到最大步数，episode结束")

        # 应用奖励归一化
        normalized_reward = self.normalize_reward(reward)
        
        # 累积当前episode的奖励（使用原始奖励进行累积）
        self.current_episode_reward += reward
        
        print(f"  -> 原始奖励: {reward:.2f}, 归一化后奖励: {normalized_reward:.2f}")
        
        obs = self.get_obs()
        return obs, normalized_reward, done, info

    def check_task_completion(self):
        """
        检查任务是否已完成
        任务完成条件：输入了文本并且保存了文件
        """
        # 如果已经给过奖励，直接返回False
        if self.task_reward_given:
            return False
            
        if not self.text_entered or not self.file_saved:
            return False
            
        # 检查目录中是否有新创建的文件
        txt_files = [f for f in os.listdir(self.text_dir) if f.endswith('.txt')]
        if not txt_files:
            return False
            
        # 检查最新文件的内容
        latest_file = max(txt_files, key=lambda f: os.path.getctime(os.path.join(self.text_dir, f)))
        try:
            with open(os.path.join(self.text_dir, latest_file), 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
                # 只要有内容就认为任务完成，不需要严格匹配特定文本
                if content.strip():  # 检查是否有非空白字符
                    print(f"  -> ******************************************")
                    print(f"  -> ******************************************")
                    print(f"  -> 检测到任务已完成: 文件 {latest_file} 包含内容")
                    print(f"  -> ******************************************")
                    print(f"  -> ******************************************")
                    # 给予任务完成奖励
                    reward = 30.0  # 增加任务完成的基础奖励
                    
                    # 给予完成任务的奖励（增加快速完成奖励）
                    speed_reward = max(0, (self.max_steps - self.step_count) * 1.0)
                    reward += speed_reward
                    print(f"  -> 奖励: 快速完成任务 (+{speed_reward:.1f})")
                    
                    # 标记已给予奖励
                    self.task_reward_given = True
                    
                    # 将奖励加到当前episode奖励中（使用原始奖励）
                    self.current_episode_reward += reward
                    
                    # 同时将奖励添加到奖励历史中用于归一化
                    self.reward_history.append(reward)
                    if len(self.reward_history) > self.reward_history_size:
                        self.reward_history.pop(0)
                    
                    print(f"  -> 任务完成，episode结束，获得奖励: {reward}")
                    return True
        except Exception as e:
            print(f"检查文件时出错: {e}")
            
        return False
    
    def is_notepad_running(self):
        """检查是否有"记事本"标题的窗口"""
        try:
            import pygetwindow as gw
            windows = gw.getWindowsWithTitle("记事本")
            return len(windows) > 0
        except:
            return False

    def open_notepad(self):
        """打开记事本"""
        # 先关闭所有记事本窗口确保干净环境
        self.close_notepad_windows()
        time.sleep(0.1)
        
        pyautogui.hotkey('win', 'r')
        pyautogui.write('notepad')
        time.sleep(0.1)  # 等待记事本启动
        pyautogui.press('enter')
        pyautogui.press('enter')
        time.sleep(0.1)  # 等待记事本启动

    def close_notepad_windows(self):
        """关闭所有记事本窗口，不触发保存对话框"""
        try:
            import pygetwindow as gw
            import subprocess
            import signal
            
            # 方法1: 使用taskkill命令强制终止记事本进程
            subprocess.run(["taskkill", "/f", "/im", "notepad.exe"], 
                        stdout=subprocess.DEVNULL, 
                        stderr=subprocess.DEVNULL)
            time.sleep(0.1)
            
        except ImportError:
            # 如果subprocess不可用，则尝试另一种方法
            try:
                # 方法2: 多次按ESC键然后强制关闭
                import pygetwindow as gw
                windows = gw.getWindowsWithTitle("记事本")
                for window in windows:
                    window.activate()
                    # 发送ESC键取消可能的对话框
                    pyautogui.press('esc')
                    time.sleep(0.05)
                    # 直接终止进程
                    window.close()
                time.sleep(0.1)
            except:
                pass
        except:
            pass

    def get_episode_rewards(self):
        """获取所有episode的奖励历史"""
        return self.episode_rewards.copy()
    
    def get_average_reward(self):
        """获取平均奖励"""
        if len(self.episode_rewards) > 0:
            return sum(self.episode_rewards) / len(self.episode_rewards)
        return 0